# Copyright (c) 2023 Presto Labs Pte. Ltd.
# Author: jingyuan

import json
import logging
import os
import shutil
import subprocess
import tempfile

import pandas as pd
from absl import app, flags
from coin2.support.product_info.product_info_holder import ProductInfoHolder

import python.coin2.api_client.kucoin as kucoin
from cc.appcoin2.strategy.support.util import get_cmd_vault_key_copy
from python.coin_deploy.basis_strat2.script.adjust_universe import get_max_pos


def update_leverage_config(config_path, leverage_map):
  with open(config_path) as f:
    config = json.load(f)
  for _, product in config['products'].items():
    symbol = product['trade_symbol'].split(':')[-1]
    leverage = min(leverage_map.get(symbol, 10), 10)
    product['execution']['leverage'] = leverage
  with open(config_path, 'w') as f:
    json.dump(config, f, indent=2)      


class MySwapClient(kucoin.KucoinFuturesPrivateClient):
  def __init__(self, key):
    super().__init__(key)
    pi_df = ProductInfoHolder().df_pi[['mea', 'symbol', 'base',
                                       'native_quote', 'native_symbol', 'contract_value']]
    pi_df = pi_df[(pi_df['mea'] == 'Futures.Kucoin.v1')
                  & (pi_df['native_quote'] == 'USDT')][['symbol', 'base', 'native_symbol', 'contract_value']]
    self.pi_df = pi_df

  def get_leverage(self):
    response = self._query.query(
        method="GET", path="/api/v1/positions", params={})
    df = pd.DataFrame(response.json()['data'])
    columns = "symbol,autoDeposit,riskLimit,currentQty,markValue,realLeverage,avgEntryPrice,liquidationPrice".split(
        ',')
    self.df = df.set_index('symbol')
    df = df[columns]
    return df

  def list_all_symbols(self):
    return self.pi_df['native_symbol'].unique()

  def lookup_symbol(self, symbol, return_type):
    assert return_type in ['native_symbol', 'symbol']
    if symbol in self.pi_df['symbol'].unique():
      symbol = self.pi_df.loc[self.pi_df['symbol'] == symbol][return_type].values[0]
    if symbol in self.pi_df['base'].unique():
      symbol = self.pi_df.loc[self.pi_df['base'] == symbol][return_type].values[0]
    if symbol in self.pi_df['native_symbol'].unique():
      symbol = self.pi_df.loc[self.pi_df['native_symbol'] == symbol][return_type].values[0]
    return symbol

  def set_leverage(self, symbols, max_pos, auto_deposit=True, config_path=None):
    leverage_map = {}
    for symbol in symbols:
      symbol = self.lookup_symbol(symbol, 'native_symbol')
      norm_symbol = self.lookup_symbol(symbol, 'symbol')

      all_limits = self._query.query(
          method="GET", path=f'/api/v1/contracts/risk-limit/{symbol}', params={}).json()['data']
      df = pd.DataFrame(all_limits)
      try:
        level = df[df['maxRiskLimit'] >= max_pos].sort_values('maxRiskLimit').iloc[0]['level']
      except IndexError:
        level = df['level'].max()
        print(
            f"{symbol} cannot satisfy {max_pos}, using max {df['maxRiskLimit'].max()} level {level}")
        print(df)
      leverage = df[df['level'] == level]['maxLeverage'].iloc[0]
      leverage_map[norm_symbol] = int(leverage)
      if symbol in self.df.index and self.df.loc[symbol]['riskLimit'] >= max_pos:
        row = self.df.loc[symbol]
        print(f"riskLimit {row['riskLimit']} >= max_pos {max_pos} for {symbol}, skip")
      else:
        resp = self._query.query(method="POST", path="/api/v1/position/risk-limit-level/change",
                                 params={}, json={'symbol': symbol, 'level': int(level)})
        print(f"updating {symbol} risk level to {level} {resp.text}")

      if symbol in self.df.index and self.df.loc[symbol]['autoDeposit'] == auto_deposit:
        row = self.df.loc[symbol]
        print(f"autoDeposit {row['autoDeposit']} == {auto_deposit} for {symbol}, skip")
      else:
        print({'symbol': symbol, 'status': auto_deposit})
        resp = self._query.query(method="POST", path="/api/v1/position/margin/auto-deposit-status",
                                 params={}, json={'symbol': symbol, 'status': auto_deposit})
        print(f"updating {symbol} auto deposit to {auto_deposit} {resp.json()}")
    if config_path:
      update_leverage_config(config_path, leverage_map)


def init_flags():
  flags.DEFINE_string('strategy', '', 'strategy_name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('key_filepath', '', 'key filepath')
  flags.DEFINE_string('symbols', 'TOMO', 'native symbol')
  flags.DEFINE_string('config_path', None, 'Update leverage config')
  flags.DEFINE_integer('maxpos', None, 'maxpos in USDT')
  flags.DEFINE_bool('query_only', True, 'check leverage')
  flags.DEFINE_bool('auto_deposit', True, 'auto_deposit')
  flags.DEFINE_bool('update_all', False, 'update all symbols')
  flags.DEFINE_bool('is_fund_account', False, 'for fund account')


def get_kucoin_account(strategy):
  if os.path.exists(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'):
    driver = json.load(open(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'))
    for exch, config in driver['order'].items():
      if config['mea'] == 'Futures.Kucoin.v1':
        return config['connection_config']['key_filepath'].split('/')[3]
  raise ValueError(f"look up account for {strategy} failed")


def check_max_pos(msc, max_pos):
  products = []
  config_path = f'python/coin_deploy/basis_strat2/config/{flags.FLAGS.strategy}.json'
  if not os.path.exists(config_path):
    return
  with open(config_path) as fi:
    config = json.load(fi)
  for key in config['products']:
    if config['products'][key].get('trade_symbol', False):
      products.append(config['products'][key]['trade_symbol'])
    elif config['products'][key].get('symbol', False):
      products.append(config['products'][key]['symbol'])
  products = [symbol.split(':')[-1] for symbol in products]

  df = msc.df
  for symbol in products:
    symbol = msc.lookup_symbol(symbol, 'native_symbol')
    if symbol not in df.index:
      continue
    row = df.loc[symbol]
    max_notional_value = float(row['riskLimit'])
    if max_pos > max_notional_value:
      logging.warning(
          f'{flags.FLAGS.strategy} {symbol} max_allowed_postion {max_notional_value} is smaller than the current max position {max_pos}')


def main(_):
  pd.set_option('display.max_rows', None)
  if flags.FLAGS.strategy and not flags.FLAGS.account:
    account = get_kucoin_account(flags.FLAGS.strategy)
    print(f'{flags.FLAGS.strategy} is using account {account}')
  else:
    account = flags.FLAGS.account
  tmp_dir = tempfile.mkdtemp()

  try:
    if flags.FLAGS.key_filepath:
      key_file = flags.FLAGS.key_filepath
    else:
      key_file = os.path.join(tmp_dir, 'key.json')
      vault_key_name = f"trade/{account}/kucoin_fut/trade_key.json"
      cmd = get_cmd_vault_key_copy(vault_key_name, key_file, flags.FLAGS.is_fund_account)
      output = subprocess.check_output(cmd, shell=True)

    msc = MySwapClient(key_file)

    results = msc.get_leverage()
    print("\n#####################")
    print("-- Before Adjust --")
    print(results)
    print("######################\n")
    if not flags.FLAGS.query_only:
      if flags.FLAGS.update_all:
        symbols = msc.list_all_symbols()
        flags.FLAGS.symbols = ",".join(symbols)
      assert flags.FLAGS.symbols

      msc.set_leverage(flags.FLAGS.symbols.split(','), max_pos=flags.FLAGS.maxpos,
                       auto_deposit=flags.FLAGS.auto_deposit,
                       config_path=flags.FLAGS.config_path)
      print("\n#####################")
      print("-- After Adjust --")
      res = msc.get_leverage()
      print(res)
      print("######################\n")

      max_pos = get_max_pos(flags.FLAGS.strategy)
      if max_pos:
        check_max_pos(msc, max_pos)
  except BaseException:
    raise
  finally:
    shutil.rmtree(tmp_dir)


if __name__ == '__main__':
  init_flags()
  app.run(main)

# Usage
# ./pyrunner cc/appcoin2/strategy/support/kucoin_leverage.py --strategy=model_smm_kucoin_perpetual_1 --query_only=0 --maxpos=18000 --symbol=BTC
# auto_deposit: kucoin only has isolated margin, enable this or you can get liqudated
# ref: https://www.kucoin.com/futures/refer/marginTrade
# Update all symbols:
# ./pyrunner cc/appcoin2/strategy/support/kucoin_leverage.py --strategy=model_smm_kucoin_perpetual_1 --query_only=0 --maxpos=18000 --update_all